import os
import json
from vector_search_hnsw import OllamaVectorSearch

# 数据库连接参数
db_params = {
    "dbname": "pgvector",
    "user": "pgvector",
    "password": "pgvector",
    "host": "10.48.0.81",
    "port": 54333
}

# 初始化搜索器
searcher = OllamaVectorSearch(db_params)

try:
    # 初始化数据库（如果需要）
    searcher.init_db()
    
    # 获取data目录下所有json文件
    data_dir = os.path.join(os.path.dirname(__file__), 'data')
    json_files = [f for f in os.listdir(data_dir) if f.endswith('.json')]
    
    for json_file in json_files:

        file_path = os.path.join(data_dir, json_file)
        
        # 读取json文件
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)


        # 获取embedding
        embedding = data.get('embedding')
        if embedding:
            # 使用文件名作为文本内容
            text = json_file
            #print(f"{embedding}================{text}========")
            # # 直接插入数据库
            with searcher.conn.cursor() as cur:
                cur.execute(
                    "INSERT INTO hnsw_embedding_documents (content, embedding) VALUES (%s, %s::halfvec)",
                    (text, embedding)
                )
            searcher.conn.commit()
            print(f"已插入文件: {json_file}")
        else:
            print(f"文件 {json_file} 中未找到embedding数据")

finally:
    searcher.close()